%% Solve standard form model using AMPL
% Simulate with observed supply/demand shocks

% Updated on 2021/03/13, Code by Sihwan Yang
%% AMPL Setting
clear
clc
tic

setting = 1; % David : 1, Sihwan : 2
adj = 10; % adjustment : 5, upperbound : 10

theta_sr = [0.7, 0.001, 0.3, 0.2];
theta_mr = [0.95, 0.001, 0.5, 0.5];

if setting == 1
% David
try
run('C:\Program Files\Artelys\AMPL 11.0.20180624\amplapi\examples\matlab\setupOnce.m');
end

else
% Sihwan
addpath('/Users/sihwan/Documents/MATLAB/Baqaee_RA/Darwinian_Frontier_AMPL/amplapi/examples/matlab');
end
%% Load shocks
N = 66; % Number of BEA sectors

BLS_shock_raw = readmatrix('BLS_labor_shock_August_All.xls','Sheet','Sheet1');
PCE_shock_raw = readmatrix('PCE_shock_August_All.xls','Sheet','Sheet1');

BLS_shock_raw(isnan(BLS_shock_raw))=0;
PCE_shock_raw(isnan(PCE_shock_raw))=0;

BLS_shock_raw = BLS_shock_raw(1:N,2:end); % Mar, Apr, May, Jun, Jul
PCE_shock_raw = PCE_shock_raw(1:N,2:end); % Mar, Apr, May, Jun

through = 2; % Apr
peak = 4; % Jun

%% Data: IO matrix(Omega), expenditure share(beta)

% Import IO data matrix
if (exist('IO_data_2018.mat', 'file') ~= 2)
    count = 1;
for year = [(1997:2018) ]
        [~, ~, raw] = xlsread('IOUse_Before_Redefinitions_PRO_1997-2015_Summary+2018.xlsx',num2str(year),'A8:CV94');
        R = cellfun(@(x) ~isnumeric(x) && ~islogical(x),raw); % Find non-numeric cells
        raw(R) = {NaN}; % Replace non-numeric cells
        Data_raw(:,:,count)= reshape([raw{:}],size(raw));
        count = count+1;
end
    [~, ~, indname] = xlsread('IOUse_Before_Redefinitions_PRO_1997-2015_Summary+2018.xlsx','1997','B8:B90');
    indname(cellfun(@(x) ~isempty(x) && isnumeric(x) && isnan(x),indname)) = {''};

    clear ans count R raw year
    save IO_data_2018.mat
else
    load IO_data_2018.mat
end
start_date = 1997;

% Calibrate the IO Table and build Omega matrix (without markups)

year = 2015; % year for which to run the calibration. 

IO = Data_raw;
Final = IO(:,99,:);
IO(:,1:2,:) = [];
IO(isnan(IO)) = 0;
IO(IO < 0) = 0;
grossout = squeeze(sum(IO(1:66,1:66,:)))+squeeze(IO(77,1:66,:))+squeeze(IO(79,1:66,:));
labor=(squeeze(IO(77,1:66,:)));
gos=(squeeze(IO(79,1:66,:))); %gross operating surplus

IO(67:end,:,:) = [];
IO(:,67:end,:) = [];

temp = IO(:,:,year-1996);
temp = temp';
Omega = bsxfun(@rdivide, temp, sum(temp,2)); 
K= gos(:,year-1996);
L= labor(:,year-1996);
va_sales=(L+ K);
alphaL=L./va_sales;
alphaK=K./va_sales;
%alphaK = max(alphaK,0); 
N = length(Omega);
beta = Final(1:N,year-1996);
beta(beta<0) = 0; %remove industries with negative implied final sales
beta = beta/sum(beta);
    zero= find(beta==0);

int_sales=squeeze(sum(IO(1:66,1:66,year-1996)))';
int_share=int_sales./(int_sales+va_sales);
va_share=(L+K)./(int_sales+va_sales);
average_mu = 1; 
mu = 1; 

%% Write Omega in standard form - relabeling
%Omega_re= final consumers + N final produc.+ N va. produc. + N int. produc. +  2FACTORS  

labor_temp = labor./grossout;
gos_temp =gos./grossout; %gross operating surplus
alphaL= labor_temp(:,year-1996);
mu = min(mu, 1./alphaL); 
alphaK= gos_temp(:,year-1996) - 1 +1./average_mu';
alphaK = max(alphaK,0); 


F=2*N;
D=1+3*N+F+3; 
% 1 is consumption today, 2:1+N is goods today, N+2:2*N+1 is VA today,
% 2*N+2:3*N+1 is intermediates today, 3*N+2:4*N+1 is labor today, 
%4*N+2:5*N is capital today, 5*N+1 is HtM consumer, 5*N+2 is the Ricardian
%consumer, 5*N+3 is (aggregate) consumption good tomorrow. 

Omega_re= zeros(D,D);
Omega_re(1,2:N+1)= beta;
Omega_re(2:1+N,1+N+1:2*N+1) = diag(mu.^(-1))*diag(va_share);
Omega_re(2:1+N,1+2*N+1:3*N+1) = diag(mu.^(-1))*diag(int_share);
for x = 1:N
Omega_re(1+N+x, 1+3*N+x)=alphaL(x)/(alphaL(x)+alphaK(x));
Omega_re(1+N+x, 1+4*N+x)=alphaK(x)/(alphaL(x)+alphaK(x));
end
Omega_re(1+2*N+1:1+3*N,2:N+1) = Omega;

Omega_re(1+5*N+1,1)=1;
Omega_re(1+5*N+2,1)=1/2;
Omega_re(1+5*N+2,end)=1/2; 

one=ones(1,1);
mu_ones= ones(D,1);
in_mu=mu_ones;
Omega_tilde= diag(in_mu)*Omega_re;
   
Psi_re = inv(eye(D)-Omega_re);
Psi_tilde = inv(eye(D)-Omega_tilde);

%% Generate Demand shock

if setting == 1
% David
try
run('C:\Program Files\Artelys\AMPL 11.0.20180624\amplapi\examples\matlab\setupOnce.m');
end 

else
% Sihwan
addpath('/Users/sihwan/Documents/MATLAB/Baqaee_RA/Darwinian_Frontier_AMPL/amplapi/examples/matlab');
end


frB_stack = zeros(3,D); % Does not depend on timing of shocks
beta_new_stack = zeros(N,1);

for i=1:3 % Elasticity (Short run, Medium run, Cobb-Douglas)

% (1) Choice of Elasticity
if i == 1 % Short run
    sigma = theta_sr(1);
    theta1 = theta_sr(2)*ones(N,1);
    epsilon = theta_sr(3)*ones(N,1);
    eta = theta_sr(4)*ones(N,1);
elseif i == 2 % Medium run
    sigma = theta_mr(1);
    theta1 = theta_mr(2)*ones(N,1);
    epsilon = theta_mr(3)*ones(N,1);
    eta = theta_mr(4)*ones(N,1);
elseif i == 3 % Cobb-Douglas
    sigma = 0.99; 
    theta1 = 0.99*ones(N,1); 
    epsilon = 0.99*ones(N,1);
    eta = 0.99*ones(N,1); 
end

% (2) Choice of Period
    shock_supply = -BLS_shock_raw(:,through+adj); % adjustment(+5), upperbound(+10)
    shock_demand = -PCE_shock_raw(:,through); 

beta_new = beta.*(1-shock_demand);
beta_new = beta_new/sum(beta_new);
beta_new_stack(:,1) = beta_new;

beta_new_std = zeros(D,1);
beta_new_std(2:N+1) = beta_new;

factor_ela=.5*ones(F,1);
rho = .99;
theta = cat(1,sigma,epsilon,eta, theta1,factor_ela,.9,rho,.9);

%Create a categorical variable to set if conditions in constraints in AMPL
factor=ones(D,1); %one for goods
factor(1) = -1; % variable part
factor(1+3*N+1:1+5*N, 1) = 0; % zero for factors
factor(end) = 0; % Since consumption tomorrow is treated like a factor. 
factor(end-2) = 3; % since the second and third to last rows are consumers. 
factor(end-1) = 2; % since the second and third to last rows are consumers. 

keynes = ones(D,1);  % zero means flexible, -1 means sticky. 
keynes(1+4*N+1:1+5*N, 1) = 0; %zero for flexible factors: capital
keynes(1+5*N+3, 1) = 0; % zero for tomorrow's consumption
keynes(1+3*N+1:1+4*N, 1) = 0; %-1 negative one for sticky factors: labor.

variab = zeros(D,1);
variab(1) = -1; % first row
variab(2:N+1) = 1; % treat this as variable

temp_factor = factor; 
%clearvars -except factor lambdaL lambdaK in_mu mu Ind F D N A L Omega_re theta IO Sigma va_share int_share alphaL alphaK

%assign ownership of the factors

chi = (factor==0).*rand(D,1);
chi(end) = 0; % tomorrow's consumption is spent by Ricardian household. 
chi(1+4*N+1:1+5*N, 1)=0; % HtM don't own capital
chi(1+3*N+1:1+4*N, 1)=0; % HtM don't own capital
 
% Determin price elasticity of factor supply 
phi = zeros(D,1); 
phi(1+4*N+1:1+5*N, 1) = 0;

%run setUponce file in case Matlab does not connect to AMPL
%'C:\Users\pc\Desktop\amplide.mswin64\amplapi\examples\matlab\setupOnce.m'
 
 %%Change Ind format for set
Ind= (1:D)';
Ind = string(Ind);           
% Ind([1 2 3 4 5 6 7 8 9], : ) = ["01","02","03","04","05","06","07","08", "09"]    
Ind = cellstr(Ind);
Indu=Ind;

Inda = (1:1)';
Inda = string(Inda);
Inda = cellstr(Inda);

Domar = Psi_re(1,2:N+1);
h = 10^(-2);

init_lambda = Psi_re(1,:)';
init_lambda(1) = 1;
init_lambda(end-2) = (Psi_re(1,:)*chi)*init_lambda(1);
init_lambda(end-1) = 2*(1-(Psi_re(1,:)*chi))*init_lambda(1);
init_lambda(end) = 0.5*init_lambda(end-1); 

init_lambda1 = init_lambda;
init_p = ones(D,1);
init_frB = ones(D,1);

Index = zeros(1,D);
foo = sort(Ind);
for v = 1:D
    Index(1,v) = str2num(foo{v});
end
foo = sortrows([Index' (1:D)']);
Index = foo(:,2);

tth=0.1;

for tt=0:tth:1

if setting == 1
% David
     try
ampl.close()
     end
ampl = AMPL('C:\Program Files\Artelys\AMPL 11.0.20180624');
ampl.setOption('solver', 'knitroampl')

else
% Sihwan
     try
ampl.close()
end
setupOnce
ampl = AMPL('/Users/sihwan/Documents/MATLAB/AMPL_Mac');
ampl.setOption('solver', 'knitro') 
end

A = ones(D,1);
A(1+3*N+1:1+4*N) =(1-tt*shock_supply); % Supply shock fully materialized

B = ones(D,D);
B(1,:)=(1+0*rand(D,1))'.*(Omega_re(1,:)>0);
B(1,2:N+1)=(1-tt*shock_demand);
B(1,:) = B(1,:)/sum(Omega_re(1,:)*B(1,:)');

in_mu=ones(D,1);
mu=ones(D,1);

ampl.read ('standard_form_covid_HtM2_demand_shock.mod')

%%Send data to ampl 
%%Choose Omega_re/Omega_tilde and Psi_re/Psi_tilde.

dfa = DataFrame(1, 'Inda');
dfa.setColumn('Inda',Inda);
ampl.setData(dfa,'Inda');

df = DataFrame(2, 'Ind', 'Indu', 'Omega');
df.setMatrix(Omega_tilde, Ind, Indu)
ampl.eval('set S1; set S2; param TheParameter{S1, S2};');
ampl.getSet('S1').setValues(Ind);
ampl.getSet('Indu').setValues(Indu);
ampl.setData(df);       

df = DataFrame(2, 'Ind', 'Indu', 'B');
df.setMatrix(B, Ind, Indu)
ampl.eval('set S11; set S12; param TheParameter1{S11, S12};');
ampl.getSet('S11').setValues(Ind);
ampl.getSet('Indu').setValues(Indu);
ampl.setData(df);  

 df1 = DataFrame(1, 'Ind','A', 'mu', 'theta','factor', 'in_mu','init_p','init_lambda','init_lambda1','keynes','chi','phi','init_frB','beta_new','variab');
    df1.setColumn('Ind', Ind );
    df1.setColumn('A', A );
    df1.setColumn('mu', mu );
    df1.setColumn('theta', theta ); 
    df1.setColumn('factor', factor);
    df1.setColumn('in_mu', in_mu);
    df1.setColumn('init_p', init_p);
    df1.setColumn('init_lambda', init_lambda);
    df1.setColumn('init_lambda1', init_lambda1);
    df1.setColumn('keynes', keynes);
    df1.setColumn('chi',chi);
    df1.setColumn('phi',phi);
    df1.setColumn('init_frB',init_frB);
    df1.setColumn('beta_new',beta_new_std);
    df1.setColumn('variab',variab);
  ampl.setData(df1, 'Ind');
 
 %%Solve  
  ampl.solve();
  
  
ampl.eval('var return_code:=solve_result_num;')
ampl.eval('display return_code;')
return_code = ampl.getVariable('return_code');
df = return_code.getValues('val');
return_code = df.getColumnAsDoubles('return_code.val');

p = ampl.getVariable('p');
df = p.getValues('val');
p = df.getColumnAsDoubles('p.val'); 

 lambda = ampl.getVariable('lambda');
 df = lambda.getValues('val');
 lambda = df.getColumnAsDoubles('lambda.val');
 
  frB = ampl.getVariable('frB');
 df = frB.getValues('val');
 frB = df.getColumnAsDoubles('frB.val');
 
 init_p = p(Index);
 init_lambda1 = lambda(Index);
 init_frB = frB(Index);
 
end

 frB_stack(i,:) = init_frB';

end

%% Evaluate GDP, Welfare 

if setting == 1
% David
try
run('C:\Program Files\Artelys\AMPL 11.0.20180624\amplapi\examples\matlab\setupOnce.m');
end 

else
% Sihwan
addpath('/Users/sihwan/Documents/MATLAB/Baqaee_RA/Darwinian_Frontier_AMPL/amplapi/examples/matlab');
end

% Welfare 1: Demand shock

Delta_r_GDP_stack1 = zeros(3,3,3); % Elasticity/Shock Order/Intensity

for i=1:3 % Elasticity (Short run, Medium run, Cobb-Douglas)

% (1) Choice of Elasticity
if i == 1 % Short run
    sigma = theta_sr(1);
    theta1 = theta_sr(2)*ones(N,1);
    epsilon = theta_sr(3)*ones(N,1);
    eta = theta_sr(4)*ones(N,1);
elseif i == 2 % Medium run
    sigma = theta_mr(1);
    theta1 = theta_mr(2)*ones(N,1);
    epsilon = theta_mr(3)*ones(N,1);
    eta = theta_mr(4)*ones(N,1);
elseif i == 3 % Cobb-Douglas
    sigma = 0.99; 
    theta1 = 0.99*ones(N,1); 
    epsilon = 0.99*ones(N,1);
    eta = 0.99*ones(N,1); 
end

    shock_supply = -BLS_shock_raw(:,through+adj); % adjustment(+5), upperbound(+10)
    shock_demand = 1-squeeze(frB_stack(i,2:N+1));

factor_ela=.5*ones(F,1);
rho = .99;
theta = cat(1,sigma,epsilon,eta, theta1,factor_ela,.9,rho,.9);

%Create a categorical variable to set if conditions in constraints in AMPL
factor=ones(D,1); %one for goods
factor(1+3*N+1:1+5*N, 1) = 0; % zero for factors
factor(end) = 0; % Since consumption tomorrow is treated like a factor. 
factor(end-2) = 3; % since the second and third to last rows are consumers. 
factor(end-1) = 2; % since the second and third to last rows are consumers. 

keynes = ones(D,1);  % zero means flexible, -1 means sticky. 
keynes(1+4*N+1:1+5*N, 1) = 0; %zero for flexible factors: capital
keynes(1+5*N+3, 1) = 0; % zero for tomorrow's consumption
keynes(1+3*N+1:1+4*N, 1) = 0; %-1 negative one for sticky factors: labor.


temp_factor = factor; 
%clearvars -except factor lambdaL lambdaK in_mu mu Ind F D N A L Omega_re theta IO Sigma va_share int_share alphaL alphaK

%assign ownership of the factors

chi = (factor==0).*rand(D,1);
chi(end) = 0; % tomorrow's consumption is spent by Ricardian household. 
chi(1+4*N+1:1+5*N, 1)=0; % HtM don't own capital
chi(1+3*N+1:1+4*N, 1)=0; % HtM don't own capital
 
% Determin price elasticity of factor supply 
phi = zeros(D,1); 
phi(1+4*N+1:1+5*N, 1) = 0;

%run setUponce file in case Matlab does not connect to AMPL
%'C:\Users\pc\Desktop\amplide.mswin64\amplapi\examples\matlab\setupOnce.m'
 
% Change Ind format for set
Ind= (1:D)';
Ind = string(Ind);           
% Ind([1 2 3 4 5 6 7 8 9], : ) = ["01","02","03","04","05","06","07","08", "09"]    
Ind = cellstr(Ind);
Indu = cellstr(Ind);

Domar = Psi_re(1,2:N+1);
h = 10^(-2);

init_lambda = Psi_re(1,:)';
init_lambda(1) = 1;
init_lambda(end-2) = (Psi_re(1,:)*chi)*init_lambda(1);
init_lambda(end-1) = 2*(1-(Psi_re(1,:)*chi))*init_lambda(1);
init_lambda(end) = 0.5*init_lambda(end-1); 

Index = zeros(1,D);
foo = sort(Ind);
for v = 1:D
    Index(1,v) = str2num(foo{v});
end
foo = sortrows([Index' (1:D)']);
Index = foo(:,2);

        for k=1:3 % Shock Order (S+D, S first, D first)

th = 0.05;
mid = 1/th/2+1;

shock_step = 0;
Delta_r_GDP = zeros(length(0:th:1),1);

nominal_GDP_s = [];
c_shares_s = [];
prices_s = [];

    init_lambda1 = init_lambda; 
    init_p = ones(D,1);

            for t=0:th:1 % Shock Intensity (0, 0.5, 1)

shock_step = shock_step + 1;  

if setting == 1
% David
     try
ampl.close()
     end
ampl = AMPL('C:\Program Files\Artelys\AMPL 11.0.20180624');
ampl.setOption('solver', 'knitroampl')

else
% Sihwan
     try
ampl.close()
end
setupOnce
ampl = AMPL('/Users/sihwan/Documents/MATLAB/AMPL_Mac');
ampl.setOption('solver', 'knitro') 
end

% Choose A and mu
% (3) Choice of Shock Order/Intensity

A = ones(D,1);
B = ones(D,D);
B(1,:)=(1+0*rand(D,1))'.*(Omega_re(1,:)>0);

    if k == 1 % S+D simultaneously
        A(1+3*N+1:1+4*N) =(1-t*shock_supply);
        B(1,2:N+1)=(1-t*shock_demand);
        B(1,:) = B(1,:)/sum(Omega_re(1,:)*B(1,:)');
    elseif k == 2 % S first, then D
        if t<0.5 || t==0.5
            A(1+3*N+1:1+4*N) =(1-2*t*shock_supply);
            B(1,2:N+1)=1;
            B(1,:) = B(1,:)/sum(Omega_re(1,:)*B(1,:)');
        else
            A(1+3*N+1:1+4*N) =(1-shock_supply);
            B(1,2:N+1)=(1-(2*t-1)*shock_demand);
            B(1,:) = B(1,:)/sum(Omega_re(1,:)*B(1,:)');
        end
    elseif k == 3 % D first, then S
        if t<0.5 || t==0.5
            A(1+3*N+1:1+4*N) = 1;
            B(1,2:N+1)=(1-2*t*shock_demand);
            B(1,:) = B(1,:)/sum(Omega_re(1,:)*B(1,:)');
        else
            A(1+3*N+1:1+4*N) =(1-(2*t-1)*shock_supply);
            B(1,2:N+1)=(1-shock_demand);
            B(1,:) = B(1,:)/sum(Omega_re(1,:)*B(1,:)');
        end
    end
in_mu=ones(D,1);
mu=ones(D,1);

%%Select model 
      %ampl.read ('standard_form_covid.mod')
      ampl.read ('standard_form_covid_HtM2.mod')

%%Send data to ampl 
%%Choose Omega_re/Omega_tilde and Psi_re/Psi_tilde.
df = DataFrame(2, 'Ind', 'Indu', 'Omega');
df.setMatrix(Omega_tilde, Ind, Indu)
ampl.eval('set S1; set S2; param TheParameter{S1, S2};');
ampl.getSet('S1').setValues(Ind);
ampl.getSet('Indu').setValues(Indu);
ampl.setData(df);       

df = DataFrame(2, 'Ind', 'Indu', 'B');
df.setMatrix(B, Ind, Indu)
ampl.eval('set S11; set S12; param TheParameter1{S11, S12};');
ampl.getSet('S11').setValues(Ind);
ampl.getSet('Indu').setValues(Indu);
ampl.setData(df);    

 df1 = DataFrame(1, 'Ind','A', 'mu', 'theta','factor', 'in_mu','init_p','init_lambda','init_lambda1','keynes','chi','phi');
    df1.setColumn('Ind', Ind );
    df1.setColumn('A', A );
    df1.setColumn('mu', mu );
    df1.setColumn('theta', theta ); 
    df1.setColumn('factor', factor);
    df1.setColumn('in_mu', in_mu);
    df1.setColumn('init_p', init_p);
    df1.setColumn('init_lambda', init_lambda);
    df1.setColumn('init_lambda1', init_lambda1);
    df1.setColumn('keynes', keynes);
    df1.setColumn('chi',chi);
    df1.setColumn('phi',phi);
  ampl.setData(df1, 'Ind');
 
 %%Solve  
  ampl.solve();
  
% Get solutions from AMPL

ampl.eval('var return_code:=solve_result_num;')
ampl.eval('display return_code;')
return_code = ampl.getVariable('return_code');
df = return_code.getValues('val');
return_code = df.getColumnAsDoubles('return_code.val');

p = ampl.getVariable('p');
df = p.getValues('val');
p = df.getColumnAsDoubles('p.val'); 

lambda = ampl.getVariable('lambda');
df = lambda.getValues('val');
lambda = df.getColumnAsDoubles('lambda.val'); 
 
cons_share = ampl.getVariable('cons_share');
df = cons_share.getValues('val');
cons_share = df.getColumnAsDoubles('cons_share.val');

cons_share = cons_share(Index);
init_p = p(Index);
init_lambda1 = lambda(Index); 

 if return_code == 0 || return_code == 100
     nominal_GDP_s = [nominal_GDP_s, lambda(1)];
     c_shares_s = [c_shares_s, cons_share(2:67)];
     prices_s = [prices_s, init_p(2:67)];
     c_init = init_lambda(1).*Omega_re(1,:)';
     c = lambda(1)/p(1)*(B(1,:).^sigma)'.*Omega_re(1,:)'.*(p(Index)./p(1)).^(-sigma);
     
     if size(c_shares_s,2) == 1
         Delta_r_GDP(shock_step,1) = 0;
     else
        Delta_r_GDP(shock_step,1) = 1-exp(sum(diff(log(nominal_GDP_s))-sum((c_shares_s(:,1:end-1)+c_shares_s(:,2:end))/2.*(diff(log(prices_s)')'))));
     end
 else
     Delta_r_GDP(shock_step,1) = NaN;

 end
            end
       
        Delta_r_GDP_stack1(i,k,:) = [Delta_r_GDP(1), Delta_r_GDP(mid), Delta_r_GDP(end)];

        end
end
ampl.close();

GDP_pt = zeros(3,3,3);
for i=1:3
    for k=1:3
        temp2 = squeeze(1-Delta_r_GDP_stack1(i,k,:));

            GDP_pt(i,k,:) = temp2;
            

    end
end
%% Macro results in Table 1
1-GDP_pt(:,[2 3 1],3)'

%% Micro results in Table 1
%% To compute micro welfare changes use exact-hat algebra:
%micro using final prefereces (price index): 1/sum(c_shares_s(c_shares_s(:,end)>0,end).*((1./init_p(Omega_re(1,:)>0)).^((1-sigma))))^(-1/(sigma-1))
%micro using initial preferences (price index): sum(c_shares_s(c_shares_s(:,end)>0,1).*((init_p(Omega_re(1,:)>0)).^((1-sigma))))^(-1/(sigma-1))
